# 预测相关路由
import io
import os
import shutil
from pathlib import Path

from exts import db
from flask import request, jsonify, g, send_file
from flask_sse import sse
from models import User, Model, Plant, Predict
import threading
from ACRmodel.train import train_model
from ACRmodel.predict import predict_plant


def register_predict_routes(app):
    # 查询所有预测
    @app.route('/plantACR/getPredicts', methods=['GET'])
    def get_predicts():
        resultVO = {
            "code": 200,
            "message": "成功获取预测列表",
            "data": {
                "predicts": []
            }
        }
        try:
            # 按ID升序查询所有预测
            predicts = Predict.query.order_by(Predict.id.asc()).all()
            if predicts:
                for predict in predicts:
                    predict_info = {
                        "id": predict.id,
                        "user_id": predict.user_id,
                        "user_name": predict.user_name,
                        "plant_id": predict.plant_id,
                        "plant_name": predict.plant_name,
                        "model_id": predict.model_id,
                        "model_name": predict.model_name,
                        "description": predict.description,
                        "files": predict.files,
                        "chr_all": predict.chr_all,
                        "chr_count": predict.chr_count,
                        "status": predict.status,
                        "create_time": predict.create_time,
                        "update_time": predict.update_time
                    }
                    resultVO["data"]["predicts"].append(predict_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "预测列表为空"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取预测列表失败!"
            return jsonify(resultVO)

        # 查询所有预测

    @app.route('/plantACR/getMyPredicts', methods=['GET'])
    def get_my_predicts():
        resultVO = {
            "code": 200,
            "message": "成功获取预测列表",
            "data": {
                "predicts": []
            }
        }
        try:
            # 按ID升序查询所有预测
            predicts = Predict.query.filter_by(user_id=g.id).order_by(Predict.id.asc()).all()
            if predicts:
                for predict in predicts:
                    predict_info = {
                        "id": predict.id,
                        "user_id": predict.user_id,
                        "user_name": predict.user_name,
                        "plant_id": predict.plant_id,
                        "plant_name": predict.plant_name,
                        "model_id": predict.model_id,
                        "model_name": predict.model_name,
                        "description": predict.description,
                        "files": predict.files,
                        "chr_all": predict.chr_all,
                        "chr_count": predict.chr_count,
                        "status": predict.status,
                        "create_time": predict.create_time,
                        "update_time": predict.update_time
                    }
                    resultVO["data"]["predicts"].append(predict_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "预测列表为空"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取预测列表失败!"
            return jsonify(resultVO)

    # 更新预测信息
    @app.route('/plantACR/updatePredict', methods=['POST'])
    def update_predict():
        resultVO = {
            "code": 200,
            "message": "预测信息更新成功",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            predict_id = data.get('id')
            new_plant_id = data.get('plant_id')
            new_model_id = data.get('model_id')

            new_plant = Plant.query.get(new_plant_id)
            new_model = Model.query.get(new_model_id)

            if not new_plant:
                resultVO["code"] = 404
                resultVO["message"] = "植物不存在"
                return jsonify(resultVO)

            if not new_model:
                resultVO["code"] = 404
                resultVO["message"] = "模型不存在"
                return jsonify(resultVO)

            predict = Predict.query.get(predict_id)

            if predict:
                # 检查是否修改了植物或模型
                is_plant_changed = predict.plant_id != new_plant_id
                is_model_changed = predict.model_id != new_model_id

                # 更新预测信息
                predict.plant_id = data.get('plant_id', predict.plant_id)
                predict.plant_name = new_plant.name
                predict.model_id = data.get('model_id', predict.model_id)
                predict.model_name = new_model.model_name
                predict.description = data.get('description', predict.description)
                if is_plant_changed or is_model_changed:  # 如果植物或模型被修改，状态重置为"未预测"
                    predict.status = "未预测"
                db.session.commit()

                resultVO["code"] = 200
                resultVO["message"] = "预测信息更新成功"
                resultVO["data"] = {
                    "predict": {
                        "id": predict.id,
                        "plant_id": predict.plant_id,
                        "plant_name": predict.plant_name,
                        "model_id": predict.model_id,
                        "model_name": predict.model_name,
                        "description": predict.description,
                        "files": predict.files,
                        "chr_all": predict.chr_all,
                        "chr_count": predict.chr_count,
                        "status": predict.status,
                        "create_time": predict.create_time,
                        "update_time": predict.update_time
                    }
                }
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "预测记录不存在"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，更新预测信息失败!"
            return jsonify(resultVO)

    # 删除预测
    @app.route('/plantACR/deletePredict', methods=['POST'])
    def delete_predict():
        resultVO = {
            "code": 200,
            "message": "预测删除成功",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            predict_id = data.get('id')
            predict = Predict.query.get(predict_id)

            # 预测不存在
            if not predict:
                resultVO["code"] = 404
                resultVO["message"] = "预测记录不存在"
                return jsonify(resultVO)

            # 检查权限，只有管理员或用户自己可以删除预测
            if g.role != 'admin' and g.id != predict.user_id:
                resultVO["code"] = 403
                resultVO["message"] = "权限不足，无法删除该预测记录"
                return jsonify(resultVO)

            # 删除数据库预测记录
            db.session.delete(predict)
            db.session.commit()  # 提交删除操作

            # # 删除本地预处理、预测结果文件
            # data_dir = f"ACRmodel/static/predict_data/{predict.plant_name}/{predict.id}"
            # predict_dir = f"ACRmodel/static/predict_result/{predict.plant_name}/{predict.id}"
            # if os.path.exists(data_dir):
            #     shutil.rmtree(data_dir)
            #     print("本地预处理数据删除成功")
            # if os.path.exists(predict_dir):
            #     shutil.rmtree(predict_dir)
            #     print("本地预测结果删除成功")

            # 删除文件，遍历所有植物目录（因为可以切换植物进行预测）
            base_data_dir = Path(app.root_path) / "ACRmodel/static/predict_data"
            base_result_dir = Path(app.root_path) / "ACRmodel/static/predict_result"
            predict_id_str = str(predict.id)

            # 删除predict_data下的所有对应ID目录
            if base_data_dir.exists():
                for plant_dir in base_data_dir.iterdir():
                    if plant_dir.is_dir():
                        target_dir = plant_dir / predict_id_str
                        if target_dir.exists():
                            shutil.rmtree(target_dir)
                            print(f"预处理数据删除: {target_dir}")

            # 删除predict_result下的所有对应ID目录
            if base_result_dir.exists():
                for plant_dir in base_result_dir.iterdir():
                    if plant_dir.is_dir():
                        target_dir = plant_dir / predict_id_str
                        if target_dir.exists():
                            shutil.rmtree(target_dir)
                            print(f"预测结果删除: {target_dir}")

            # 删除temp文件夹中的文件
            if predict.files:
                temp_dir = Path(app.root_path) / "ACRmodel" / "static" / "temp"
                for filename in predict.files.split(','):
                    file_path = temp_dir / filename.strip()
                    if file_path.exists():
                        os.remove(file_path)
                print("temp中文件删除成功")

            resultVO["code"] = 200
            resultVO["message"] = "预测记录删除成功"
            resultVO["data"] = {
                "predict": {
                    "id": predict.id,
                }
            }
            return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，删除预测记录失败!"
            return jsonify(resultVO)

    # 创建预测任务(未预测)
    @app.route('/plantACR/createPredict', methods=['POST'])
    def create_predict():
        resultVO = {
            "code": 200,
            "message": "已创建预测任务",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            user = User.query.get(g.id)
            plant = Plant.query.get(data.get('plant_id'))
            model = Model.query.get(data.get('model_id'))
            # 检查用户 是否存在
            if not user:
                resultVO["code"] = 404
                resultVO["message"] = "登录信息异常，用户不存在"
                return jsonify(resultVO)

            # 检查植物 是否存在
            if not plant:
                resultVO["code"] = 404
                resultVO["message"] = "对应的植物在数据库中不存在：" + data.get('plant_name')
                return jsonify(resultVO)

            # 检查模型 是否存在
            if not model:
                resultVO["code"] = 404
                resultVO["message"] = "对应的模型在数据库中不存在：" + data.get('dataset_name')
                return jsonify(resultVO)

            # 新增预测记录
            new_predict = Predict(
                user_id=user.id,
                user_name=user.name,
                plant_id=plant.id,
                plant_name=plant.name,
                model_id=model.id,
                model_name=model.model_name,
                description=data.get('description'),
                files=data.get('files'),
                chr_all=data.get('chr_all'),
                chr_count=data.get('chr_count'),
                status="未预测",
            )
            db.session.add(new_predict)
            db.session.commit()
            resultVO["data"] = {
                "predict": {
                    "id": new_predict.id,
                    "plant_id": new_predict.plant_id,
                    "plant_name": new_predict.plant_name,
                    "model_id": new_predict.model_id,
                    "model_name": new_predict.model_name,
                    "description": new_predict.description,
                    "files": new_predict.files,
                    "chr_all": new_predict.chr_all,
                    "chr_count": new_predict.chr_count,
                    "status": new_predict.status,
                    "create_time": new_predict.create_time,
                    "update_time": new_predict.update_time
                }
            }
            return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            print("预测预测任务添加失败:", e)
            resultVO["code"] = 500
            resultVO["message"] = "预测任务添加失败!"
            return jsonify(resultVO)

    # 开始预测
    @app.route('/plantACR/startPredict', methods=['POST'])
    def start_predict_plant():
        resultVO = {
            "code": 200,
            "message": "预测开始",
            "data": {}
        }
        data = request.get_json()  # 从前端获取 JSON 数据
        predict_id = data.get('id')
        try:
            # 检查预测是否存在
            predict = Predict.query.get(predict_id)
            if not predict:
                resultVO["code"] = 404
                resultVO["message"] = "预测不存在"
                return jsonify(resultVO)

            # 检查权限，只有管理员或用户自己能预测
            if g.id != predict.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = "权限不足，无法进行预测"
                return jsonify(resultVO)

            # 检查是否有其他预测任务正在进行
            predicting = Predict.query.filter_by(status="预测中").first()
            if predicting:
                resultVO["code"] = 400
                resultVO["message"] = f"预测 {predicting.id} 正在预测中，请等待完成"
                return jsonify(resultVO)

            # 检查预测记录是否可以预测，预测的状态：未预测, 预测中, 预测失败, 预测成功
            if predict.status in ('未预测', '预测失败'):
                resultVO["code"] = 200
                resultVO["message"] = "开始预测"
                resultVO["data"] = {
                    "predict": {
                        "id": predict.id,
                        "plant_name": predict.plant_name,
                        "model_name": predict.model_name,
                        "status": predict.status,
                    }
                }
                # 启动后台线程执行预测
                threading.Thread(target=start_predict, args=(predict.id,)).start()
                # 启动预测任务
                predict.status = "预测中"
                db.session.commit()
                return jsonify(resultVO)
            # 预测中状态是无法发送预测请求的，返回数据异常
            else:
                resultVO["code"] = 400
                resultVO["message"] = "数据异常！"
                return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            predict = Model.query.get(predict_id)
            if predict:
                predict.status = "预测失败"
                db.session.commit()
            print("预测任务启动失败:", str(e))
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，预测失败!"
            return jsonify(resultVO)

    # 在后台线程中执行预测
    def start_predict(predict_id):
        """
            status: 未预测, 预测中, 预测失败, 预测成功
                预测完成后，更新预测的状态

            :param
                predict_id: 预测id
            :return
        """
        with app.app_context():
            predict_result = {}
            try:
                predict = Predict.query.get(predict_id)
                if predict:
                    print("启动预测...")
                    sse_predict_status(predict_id, "预测中", "开始预测")
                    # 调用预测函数启动预测
                    predict_result = predict_plant(predict.id, predict.plant_name, predict.model_id, predict.files)

                    if predict_result["status"] == "预测成功":
                        print("预测成功，预测结果:", predict_result)
                        predict.status = predict_result["status"]
                        predict.chr_all = predict_result["chr_all"]
                        predict.chr_count = predict_result["chr_count"]
                        db.session.commit()
                        sse_predict_status(predict_id, "预测成功", "预测成功！")
                    else:
                        print("预测失败，预测结果：", predict_result)
                        sse_predict_status(predict_id, "预测失败", predict_result["message"])
                else:
                    print("预测任务不存在")
                    sse_predict_status(predict_id, "预测失败", "预测任务不存在")
            except Exception as e:
                db.session.rollback()
                print("预测失败，报错：:", e)
                sse_predict_status(predict_id, "预测失败", "服务器错误，预测失败！")
            finally:
                print("finally:当前预测结束")

    # 发送SSE，并更新预测在数据库中的状态，预测的状态：未预测, 预测中, 预测失败, 公开, 私密
    def sse_predict_status(predict_id, status, sse_message=None):
        predict = Predict.query.get(predict_id)
        if predict:
            predict.status = status
            db.session.commit()
            # 这里向所有订阅了 predict_status 事件的客户端（前端）发送一条消息
            try:
                # 发送SSE消息
                sse.publish({
                    "predict_id": predict_id,
                    "status": status,
                    "message": sse_message,
                }, type='predict_status', channel='sse')
                print("SSE消息已发布")
            except Exception as e:
                print("SSE发布失败：", e)
            return "消息已发送"
        else:
            print("预测不存在")
            sse.publish({
                "predict_id": predict_id,
                "status": "预测失败",
                "message": "预测记录不存在",
            }, type='predict_status', channel='sse')
            return "消息已发送，预测记录不存在"

    # 下载预测
    @app.route('/plantACR/downloadPredict/<int:predict_id>', methods=['GET'])
    def download_predict(predict_id):
        resultVO = {
            "code": 200,
            "message": ""
        }
        try:
            predict = Predict.query.get(predict_id)
            if not predict:
                resultVO["code"] = 404
                resultVO["message"] = "预测记录不存在"
                return jsonify(resultVO)  # 虽然是json，但是前端以二进制文件Blob接受，需要进一步读取

            # 只有管理员和用户自己可以下载
            if g.id != predict.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            # 预测文件路径
            file_name = "predict.bed"
            predict_file = Path(app.root_path) / "ACRmodel" / "static" / "predict_result" / predict.plant_name / str(
                predict_id) / file_name

            # 文件不存在
            if not predict_file.exists():
                resultVO["code"] = 404
                resultVO["message"] = "文件不存在"
                return jsonify(resultVO)

            # 获取二进制流
            with open(predict_file, 'rb') as file:
                bytes_io = io.BytesIO(file.read())
            # 返回文件给前端
            return send_file(bytes_io,
                             mimetype='text/plain, application/octet-stream',
                             as_attachment=True,  # 强制浏览器以附件形式下载文件，而非直接打开
                             download_name=file_name  # 设置默认下载文件名
                             )
        except Exception as e:
            print(str(e))
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，下载失败！"
            return jsonify(resultVO)

    # 获取可视化数据
    @app.route('/plantACR/getPredictDatas/<int:predict_id>', methods=['GET'])
    def get_predict_datas(predict_id):
        resultVO = {
            "code": 200,
            "message": "成功获取可视化数据",
            "data": {
                "predictDatas": [],
                "chromosomes": []
            }
        }
        try:
            predict = Predict.query.get(predict_id)
            if not predict:
                resultVO["code"] = 404
                resultVO["message"] = "预测记录不存在"
                return jsonify(resultVO)

            # 只有管理员和用户自己可以可视化
            if g.id != predict.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            # 预测文件不存在
            predict_file = Path(app.root_path) / "ACRmodel" / "static" / "predict_result" / predict.plant_name / str(
                predict_id) / "predict.bed"
            if not predict_file.exists():
                resultVO["code"] = 404
                resultVO["message"] = "文件不存在"
                return jsonify(resultVO)

            # 读取预测文件
            chromosomes = set()  # 用于存储所有染色体号
            predict_datas = []
            with open(predict_file, 'r') as file:
                for line in file:
                    if line.strip():  # 跳过空行
                        parts = line.strip().split()
                        if len(parts) >= 5:  # 确保有5列数据
                            chr_num = parts[0]
                            start = int(parts[1])
                            end = int(parts[2])
                            prob = float(parts[3])
                            accessible = int(parts[4])

                            # 添加到数据列表
                            predict_datas.append({
                                "chr": chr_num,
                                "start": start,
                                "end": end,
                                "prob": prob,
                                "accessible": accessible
                            })
                            # 添加到染色体集合
                            chromosomes.add(chr_num)
            # 更新返回结果
            resultVO["data"]["predictDatas"] = predict_datas
            resultVO["data"]["chromosomes"] = sorted(chromosomes, key=lambda x: (not x.isdigit(), int(x) if x.isdigit() else x))
            return jsonify(resultVO)
        except Exception as e:
            print(str(e))
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取可视化数据失败！"
            return jsonify(resultVO)

    return app
